import json
import uuid
from datetime import datetime


class MessageContent:
    def __init__(self, msg_type, data):
        self.msg_type = msg_type
        self.data = data


class Message:
    def __init__(self, message_id, role, content, timestamp):
        self.content = content
        self.message_id = message_id
        self.role = role
        self.timestamp = timestamp


class ConversationMessage:
    def __init__(self, client_type, client_id, session_id, user_name, target_client_id, messages=None,
                 code="success", action='talk'):
        self.client_type = client_type
        self.client_id = client_id
        self.session_id = session_id
        self.messages = messages if messages else []
        self.code = code
        self.user_name = user_name
        self.target_client_id = target_client_id
        self.action = action

    def to_json(self):
        return json.dumps(self, default=lambda o: o.__dict__, ensure_ascii=False)

    @staticmethod
    def from_json(json_str):
        data = json.loads(json_str)
        client_type = data.get("client_type")
        client_id = data.get("client_id")
        action = data.get("action")
        session_id = data.get("session_id")
        user_name = data.get("user_name")
        target_client_id = data.get("target_client_id")
        messages = ConversationMessage._parse_messages(data.get("messages"))
        code = data.get("code")

        return ConversationMessage(client_type, client_id, session_id, user_name, target_client_id, messages, code,
                                   action)

    @staticmethod
    def _parse_content(content):
        parsed_content = []
        for content_data in content:
            msg_type = content_data.get("msg_type")
            content_msg_data = content_data.get("data")
            parsed_content.append(MessageContent(msg_type, content_msg_data))
        return parsed_content

    @staticmethod
    def _parse_messages(messages):
        parsed_messages = []
        if isinstance(messages, dict):
            role = messages.get("role")
            timestamp = messages.get("timestamp")
            message_id = messages.get("message_id")
            content = ConversationMessage._parse_content(messages.get("content", []))
            parsed_messages.append(
                {"role": role, "timestamp": timestamp, "content": content, "message_id": message_id})
            return parsed_messages
        else:
            for message_data in messages:
                role = message_data.get("role")
                timestamp = message_data.get("timestamp")
                message_id = message_data.get("message_id")
                content = ConversationMessage._parse_content(message_data.get("content", []))
                parsed_messages.append(
                    {"role": role, "timestamp": timestamp, "content": content, "message_id": message_id})
            return parsed_messages

    @staticmethod
    def format_context(messages):
        transformed_messages = []
        for message in messages:
            role = message['role']
            content = ''
            for item in message['content']:
                if item.msg_type == 'text':
                    content = item.data
                    break
            transformed_messages.append({'role': role, 'content': content})
        return transformed_messages

    @staticmethod
    def gen_msg(msg_type, client_id, client_type, message, session_id, target_client_id="", code="success",
                role="assistant",
                message_id=None, message_time=None, action="talk"):
        send_content = MessageContent(msg_type, message)
        content_arr = [send_content]
        send_message = Message(message_id if message_id else str(uuid.uuid4()).replace("-", ""), role, content_arr,
                               message_time)
        send_msg = ConversationMessage(client_type, client_id, session_id, "", target_client_id, send_message, code,
                                       action)
        return send_msg
